/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: hdds op remote update execute.
 * Author: liutianming
 * Create: 2022-11-28
 */

#ifndef OP_HDDS_REMOTE_UPDATE_EXECUTOR_H
#define OP_HDDS_REMOTE_UPDATE_EXECUTOR_H

#include <map>
#include <unordered_map>
#include <atomic>
#include "hdds_param.h"
#include "op_embedding_service_executor.h"
#include "hccl_operator.h"
#include "transport_base_pub.h"
#include "aicpu_sharder.h" // aicpu sharding
#include "auto_buffer.h"
#include "con_unique_mappers_mgr.h"

namespace hccl {
constexpr u32 DEFAULT_USERREQUIRE_VALUE = 0;
using SendStatusRecord = struct SendStatusRecordDef {
    u64 alreadySendCount;
    bool sendCompFlag;
    SendStatusRecordDef() : alreadySendCount(0), sendCompFlag(false) {}
};

class HdcsOpRemoteUpdateExecutor : public OpEmbeddingServiceExecutor {
public:
    static constexpr u32 DOUBLE = 2;

    explicit HdcsOpRemoteUpdateExecutor();
    ~HdcsOpRemoteUpdateExecutor() override;
    HcclResult ExecuteOp(HcclOpPtr opPtr) override;
    HcclResult Init(std::vector<u32> &psRankId);
    HcclResult DeInit();

private:
    template <typename T, typename TM>
    using ConUniqueMappersInfo = typename ConUniqueMappersMgr<T, TM>::ConUniqueMappersInfo;

    HcclResult ExecuteTask(HcclOperator<EmbeddingServiceParam> *opPtr);

    HcclResult TaskUpdate(OpTaskPtr &task);
    HcclResult TaskUpdateReduceSum(OpTaskPtr &task);
    HcclResult TaskUpdateKeyReduce(OpTaskPtr &task);
    HcclResult TaskSendRequest(OpTaskPtr &task);
    HcclResult TaskResetUniqueHandle(OpTaskPtr &task);
    HcclResult TaskRecvResponse(OpTaskPtr &task);

    HcclResult ReduceSum(ReduceSumTask *params);
    HcclResult VectorReduce(u32 taskIndex);

    HcclResult KeyReduce(KeyReduceTask &params);
    HcclResult KeyReduceWithoutMappingMatrix(KeyReduceTask &params);
    HcclResult ConvertPsBufferInfo(KeyReduceTask &params);
    HcclResult GetPsIdInfo(const KeyReduceTask &params);
    HcclResult ClearValueMem(const KeyReduceTask &params);
    HcclResult ReduceKeyAndGetMappingMatrix(const KeyReduceTask &task);

    HcclResult UpdateDataSend(ReduceSumTask *params);
    HcclResult Send(ReduceSumTask *params, u32 rankId, void *buffer, u64 count,
        std::vector<HcclRequest> &requestArray, HcclUserRequire &userRequire);
    HcclResult WaitSomeAll(std::vector<HcclRequest> &requestArray);
    HcclResult WaitUpdateFinish(ReduceSumTask *params);
    HcclResult RecvRespones(ReduceSumTask *params, u32 rank, u32 &signal,
        std::vector<HcclRequest> &requestArray, std::vector<s32> &recvCompCounts);

    HcclResult GetOrPrepareMappingInfo(KeyReduceTask &params);
    HcclResult FetchUniquedCounts();
    u32 CalThreadNumEachPs(u32 keyMaxNum);
    HcclResult ValueReudceSum(KeyReduceTask &params, std::function<bool(const u32&)> filtered);
    HcclResult ParralValueReudceSum(KeyReduceTask &params);

    HcclResult PrepareEnvelopesForRoce(void *buffer, u64 count,
        HcclDataType dataType, HcclEnvelope &envelope, u64 errorStatusAddrVal,
        HcclUserRequire &userRequire);

    std::unordered_map<u32, Transport *> rankTransportMap_;
    std::map<OpTaskType, std::function<HcclResult(OpTaskPtr &)>> taskExeFuncTable_;
    u32 keyMaxNum_ = 0;
    std::vector<u32> psRankId_;
    std::vector<u32> rankId2PsId_;
    std::vector<u32> psKeyNum_;
    std::vector<u64> psOffset_;
    std::unordered_map<s64, u32> realKeyMap_;
    PsEnvelopeInfo envelopeInfo_;
    EventInfo eventRegInfo_;
    u32 *psIdsMem_ = nullptr;
    bool firstMemset_;

    u32 parallelExecCpuNum_ = 1;
    std::vector<ReduceSumTask *> vectorReduceParams_;

    std::vector<std::unordered_map<s64, u32>> realKeyMapForShard_; // shard分段做去重的map
    std::vector<std::vector<u32>> psKeyNumForShard_; // index psid index threadId value key num
    std::vector<std::vector<u64>> psOffsetForShard_; // index psid index threadId value offset
    std::vector<std::vector<u32>> threadkeyNumForShard_; // 每个线程要处理的key num以及对应的key index
    u32 psSize_ = 0;

    // 去重的数据类型为s64，映射索引Matrix数据类型为s32
    ConUniqueMappersMgr<s64, s32> conUniqueMappersMgr_{};
    AutoBuffer<s32> mappingMatrix_{};
    AutoBuffer<PsBufferInfo> psOffsetAndNum_{};
};
}

#endif
